#include "patchcoreai.h"
#include "source/configmanager.h"
#include "source/rknntool.h"
#include "errordialogmanager.h"
#include <QDir>
#include <QElapsedTimer>


PatchCoreAi::PatchCoreAi()
{
    LogDebug << "Attempting to PatchCoreAi";
    initial();
    //std::string python_executable = "/home/cat/miniconda3/envs/rknn_toolkit_lite2/bin/python";
    //setenv("PYTHONHOME", python_executable.c_str(), 1);

    // py::initialize_interpreter();  // 初始化 Python 解释器
    // py::module sys = py::module::import("sys");
    // //sys.attr("executable") = python_executable;
    // py::print(sys.attr("executable"));
    // py::print(sys.attr("path"));

    canRun=false;
    imageIndex = 0;
    ctx = 0 ;
    image = cv::Mat();
}

PatchCoreAi::~PatchCoreAi()
{
    unloadModel();

}
int PatchCoreAi::initial() {
    QString modelBeloneName = ParamValue.ccnfigfileName();

    QDir modelDir(model_Path);
    if (!modelDir.exists()) {
        if (!modelDir.mkpath(".")) {
            LogError << "Failed to create model directory: " <<
                modelDir.path().toStdString() ;
            return -1;
        }
    }
    modelBeloneDir = QDir(modelDir.filePath(modelBeloneName));
    if (!modelBeloneDir.exists()) {
        if (!modelBeloneDir.mkpath(".")) {
            LogError<< "Failed to create modelBeloneName directory: " <<
                modelBeloneDir.path().toStdString() ;
            return -1;
        }
    }

    // 获取 modelBeloneDir 目录下所有的 .jpg 文件路径
    QStringList filters;
    filters << "*.jpg";
    LogDebug << "Attempting to find jpg in " << modelBeloneDir.path().toStdString();
    QStringList imageList = modelBeloneDir.entryList(filters, QDir::Files);

    QList<QString> imagePaths;
    QRegExp rx("(\\d+)_([RD]).jpg");
    imageInfoList.clear();

    foreach (const QString &fileName, imageList) {
        QString fullPath = modelBeloneDir.filePath(fileName);
        imagePaths.append(fullPath);
        if (rx.indexIn(fileName) != -1) {
            int number = rx.cap(1).toInt();
            QString status = rx.cap(2);
            bool isTrained = (status == "R");
            QVariantMap imageInfo;
            imageInfo["istrain"] = isTrained;
            imageInfo["img"] = fullPath;
            imageInfoList.append(imageInfo);

            if (number > imageIndex) {
                imageIndex = number;
            }
        }
    }
    // LogDebug << "patchCoreAi_.imageInfoList()" << imageInfoList;
    // 在此处可以对 imageInfoList 进行处理或输出
    // 例如打印所有图片的信息
    // for (const QVariant &imageInfo : imageInfoList) {
    //     QVariantMap imageInfoMap = imageInfo.toMap();
    //     LogError << "Is Trained: " << imageInfoMap["istrain"].toBool()
    //               << ", Path: " << imageInfoMap["img"].toString().toStdString();
    // }

    // 返回成功标志
    return 0 ;
}

bool PatchCoreAi::loadModel(const QString &modelPath)
{
    canRun = false;
    if(ctx)
    {
        LogWarning << "The model has been opened and is currently being closed";
        unloadModel();
    }
    // 打开模型文件
    FILE *fp = fopen(modelPath.toStdString().c_str(), "rb");
    if (fp == nullptr) {
        ErrorDialogManager::instance().showNonBlockingError("载入模型错误", "文件不存在....");
        LogWarning << "Open file" << modelPath << "failed.";
        return false;
    }

    // 读取模型文件大小
    fseek(fp, 0, SEEK_END);
    modelDataSize = ftell(fp);
    fseek(fp, 0, SEEK_SET);

    modelDetal["size"] = QVariant::fromValue(modelDataSize);
    // 将模型载入内存
    modelData = (unsigned char*)malloc(modelDataSize);
    if (modelData == nullptr) {
        ErrorDialogManager::instance().showNonBlockingError("载入模型错误", "内存分配错误");
        LogWarning << "Malloc model data failed.";
        fclose(fp);
        return false;
    }
    fread(modelData, 1, modelDataSize, fp);
    fclose(fp);
    LogDebug << "read loadModel Successfully modelDataSize:" << modelDataSize;

    // 初始化模型
    int ret = rknn_init(&ctx, modelData, modelDataSize, 0, nullptr);
    if (ret < 0) {
        ErrorDialogManager::instance().showNonBlockingError("载入模型错误", "模型无法初始化检查文件有效性....");
        LogWarning << "rknn_init error :" << ret;
        free(modelData);
        return false;
    }

    // 获取SDK版本号
    ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version));
    if (ret < 0) {
        ErrorDialogManager::instance().showNonBlockingError("载入模型错误", "模型无法初始化，失败，检查文件有效性....");
        LogWarning << "rknn_query RKNN_QUERY_SDK_VERSION error ret: " << ret;
        return false;
    }
    LogDebug << "sdk version: " << version.api_version << " driver version: " << version.drv_version;

    // 查看模型输入输出
    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &ioNum, sizeof(ioNum));
    if (ret != 0) {
        ErrorDialogManager::instance().showNonBlockingError("载入模型错误", "模型输入输出存在问题....");
        LogWarning << "rknn_query RKNN_QUERY_IN_OUT_NUM error" << ret;
        return false;
    }
    LogDebug << "model input num: " << ioNum.n_input << " output num: " << ioNum.n_output;

    // 获取输入张量属性

    memset(inputAttrs, 0, sizeof(inputAttrs));
    input_num = ioNum.n_input;
    for (int i = 0; i < ioNum.n_input; i++) {
        inputAttrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(inputAttrs[i]), sizeof(rknn_tensor_attr));
        if (ret != 0) {
            ErrorDialogManager::instance().showNonBlockingError("载入模型错误", "无法读取输入输出....");
            LogWarning << "rknn_query RKNN_QUERY_INPUT_ATTR error" << ret;
            return false;
        }
        dump_tensor_attr(&(inputAttrs[i]));
    }

    modelDetal["type"] = QString::fromStdString(std::string(get_type_string(inputAttrs[0].type)));
    modelDetal["coreSet"] = QVariant::fromValue(inputAttrs[1].dims[0]);

    // 获取输出张量属性
    memset(outputAttrs, 0, sizeof(outputAttrs));
    for (int i = 0; i < ioNum.n_output; i++) {
        outputAttrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(outputAttrs[i]), sizeof(rknn_tensor_attr));
        if (ret != 0) {
            ErrorDialogManager::instance().showNonBlockingError("载入模型错误", "模型输出存在问题....");
            LogWarning << "rknn_query RKNN_QUERY_OUTPUT_ATTR error" << ret;
            return false;
        }
        dump_tensor_attr(&(outputAttrs[i]));
    }

    for (uint32_t i = 0; i < ioNum.n_output; ++i)
    {
        int output_size = outputAttrs[i].n_elems * 2;
        output_mems[i] = rknn_create_mem(ctx, output_size);
        if (output_mems[i] == nullptr) {
            LogError << "rknn_create_mem fail for output tensor " << i;
            ErrorDialogManager::instance().showNonBlockingError("分配输出内存错误", "无法为输出张量分配内存。");
            return false;
        }
    }

    rknn_custom_string custom_string;
    ret = rknn_query(ctx, RKNN_QUERY_CUSTOM_STRING, &custom_string, sizeof(custom_string));
    if (ret != RKNN_SUCC)
    {
        LogError << "rknn_query fail! ret=" << ret;
        ErrorDialogManager::instance().showNonBlockingError("载入模型错误", "读取自定义失败....");
        return false;
    }

    inputAttrs[0].type = RKNN_TENSOR_FLOAT16;
    inputAttrs[1].type = RKNN_TENSOR_FLOAT16;
    //   // default fmt is NHWC, npu only support NHWC in zero copy mode
    inputAttrs[0].fmt = RKNN_TENSOR_NHWC;
    inputAttrs[1].fmt = RKNN_TENSOR_UNDEFINED;

    input_mems[0] = rknn_create_mem(ctx, inputAttrs[0].size_with_stride);
    input_mems[1] = rknn_create_mem(ctx, inputAttrs[1].size_with_stride);
    if (input_mems[0] == nullptr || input_mems[1] == nullptr) {
        LogError << "rknn_create_mem fail for output tensor ";
        ErrorDialogManager::instance().showNonBlockingError("分配输出内存错误", "无法为输出张量分配内存。");
        return false;
    }

    input_size = inputAttrs[0].dims[2];
    return true;
}

bool PatchCoreAi::unloadModel()
{
    canRun=false;
    readyRun = false;
    LogInfo << "Attempting to PatchCoreAi::unloadModel";
    // if(ctx==0)
    // {
    //     LogDebug << "fail to PatchCoreAi::unloadModel because model has unload";
    //     return true;
    // }
    // LogDebug << "Attempting to unloadModel :input_mems1";
    // rknn_destroy_mem(ctx, input_mems[1]);
    // LogDebug << "Attempting to unloadModel :input_mems0";
    // //rknn_destroy_mem(ctx, input_mems[0]);

    // LogDebug << "Attempting to unloadModel :output_mems";
    // for (uint32_t i = 0; i < 2; ++i) {
    //     rknn_destroy_mem(ctx, output_mems[i]);
    // }
    // LogDebug << "Attempting to unloadModel :rknn_destroy";
    // // destroy
    // rknn_destroy(ctx);
    // LogDebug << "Attempting to unloadModel :modelData";
    // if (modelData != nullptr) {
    //     free(modelData);
    // }
    // LogInfo << "succesfull to unloadModel :modelData";
    return true;
}

bool PatchCoreAi::loadCoreSet()
{
    canRun = false;
    LogDebug << "Attempting to find coreset";
    int coreSet = modelDetal.value("coreSet").toInt();

    // 生成文件名
    QString fileName = QString("coreset_%1.bin").arg(coreSet);

    // 构造完整路径
    QString coresetFilePath = modelBeloneDir.filePath(fileName);

    if (!QFile::exists(coresetFilePath))
    {
        LogError << "File does not exist! ::" << coresetFilePath;
        //ErrorDialogManager::instance().showNonBlockingError("载入 CoreSet 错误", "文件不存在....");
        return false;
    }

    LogDebug << "Attempting to loadCoreSet :" <<coresetFilePath;
    std::ifstream file(coresetFilePath.toStdString(), std::ios::binary);

    if (!file) {
        ErrorDialogManager::instance().showNonBlockingError("载入 CoreSet 错误", "文件不存在....");
        // LogWarning << "Open file:" << coresetFilePath << "failed.";
        return false;
    }

    file.seekg(0, std::ios::end);
    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);

    if (size == 0) {
        ErrorDialogManager::instance().showNonBlockingError("载入 CoreSet 错误", "文件大小为0....");
        LogWarning << "File size is zero: " << coresetFilePath;
        file.close();
        return false;
    }
    coreSetDetal["size"] = QVariant::fromValue(size);
    coreSetDetal["c"] = QVariant::fromValue((int)(size/2/coresetDim));


    std::vector<char> buffer(size);
    file.read(buffer.data(), size);
    file.close();

    if(coreSetDetal["c"].toInt() != modelDetal["coreSet"].toInt())
    {
        ErrorDialogManager::instance().showNonBlockingError("知识库不匹配模型", "请重新加载模型或者重新训练");
        LogWarning << "model need dim: " << modelDetal["coreSet"].toInt() << " but corSet Dim is "<< coreSetDetal["c"].toInt() ;
        return false;
    }

    if (!file) {
        ErrorDialogManager::instance().showNonBlockingError("载入 CoreSet 错误", "文件读取错误....");
        LogWarning << "Failed to read file: " << coresetFilePath;
        return false;
    }

    // Assuming input_mems[1] is correctly allocated and initialized

    // memcpy(input_mems[0]->virt_addr, image.data, 0);
    memcpy(input_mems[1]->virt_addr, buffer.data(), size);

    // 设置输入张量内存
    // int ret = rknn_set_io_mem(ctx, input_mems[0], &input_attrs[0]);
    // if (ret < 0)
    // {
    //     LogError << "rknn_set_io_mem fail for input tensor 0, ret=" << ret;
    //     ErrorDialogManager::instance().showNonBlockingError("设置输入内存错误", "无法设置输入张量内存。");
    //     return false;
    // }

    int ret = rknn_set_io_mem(ctx, input_mems[0], &inputAttrs[0]);
    if (ret < 0)
    {
        LogWarning <<  "rknn_set_io_mem1 fail! ret=" << ret;
        ErrorDialogManager::instance().showNonBlockingError("设置输入内存错误", "无法设置输入张量内存。");
        return false;
    }

    ret = rknn_set_io_mem(ctx, input_mems[1], &inputAttrs[1]);
    if (ret < 0)
    {
        LogError << "rknn_set_io_mem fail for input tensor 1, ret=" << ret;
        ErrorDialogManager::instance().showNonBlockingError("设置输入内存错误", "无法设置输入张量内存。");
        return false;
    }

    // 设置输出张量内存
    for (uint32_t i = 0; i < ioNum.n_output; ++i)
    {
        // 默认输出类型取决于模型，这里需要 float32 来计算 top5
        outputAttrs[i].type = RKNN_TENSOR_FLOAT16;
        // 设置输出内存和属性
        ret = rknn_set_io_mem(ctx, output_mems[i], &outputAttrs[i]);
        if (ret < 0)
        {
            LogError << "rknn_set_io_mem fail for output tensor " << i << ", ret=" << ret;
            ErrorDialogManager::instance().showNonBlockingError("设置输出内存错误", "无法设置输出张量内存。");
            return false;
        }
    }
    ErrorDialogManager::instance().showNonBlockingInfo("载入知识库成功", "可以正常使用");
    LogDebug << "CoreSet loaded successfully from " << coresetFilePath;
    return true;
}

int PatchCoreAi::detect(cv::Mat &Image,object_detect_result_list *od_results , float objThresh)
{
    //cv::cvtColor(Image, Image, cv::COLOR_BGR2RGB);
    LogTrack << "Attempting to PatchCoreAi detect";
    cv::Size srcSize = Image.size();
    cv::Mat resized_image;
    cv::resize(Image, resized_image, cv::Size(input_size, input_size), 0, 0, cv::INTER_LINEAR);
    cv::Mat fp16_image(resized_image.size(), CV_16FC3);
    if (modelDetal["type"].toString() == "FP16")
    {
        resized_image.convertTo(fp16_image, CV_16FC3);
        if(inputAttrs[0].size_with_stride != fp16_image.total() * fp16_image.elemSize())
        {
            LogWarning << "inputAttrs[0].size_with_stride != fp16_image.total() * fp16_image.elemSize()" ;
            return -1;
        }
        memcpy(input_mems[0]->virt_addr, fp16_image.data, inputAttrs[0].size_with_stride);
        LogTrack << "Attempting to Prknn_run input_size:"<< input_size ;
        QElapsedTimer timer;
        timer.start();

        int ret = rknn_run(ctx, NULL);
        if (ret < 0)
        {
            LogWarning << "rknn run error: " << ret;
            return -1;
        }
        qint64 elapsed = timer.nsecsElapsed(); // 获取执行时间（纳秒）
        LogTrackq << "rknn_run 函数执行时间:" << elapsed << "纳秒";


        fp16_t *fp16_ptr = reinterpret_cast<fp16_t *>(output_mems[1]->virt_addr);
        cv::Mat fp16_mat(1, 1, CV_16FC1, fp16_ptr);
        // 将FP16值转换为float
        cv::Mat float_mat;
        fp16_mat.convertTo(float_mat, CV_32F);

        // 读取转换后的float值
        float soure_score = float_mat.at<float>(0, 0);
        pred_score_ = std::sqrt(soure_score);
        LogTrack << "Patch Core pred_score :" << pred_score_;


        cv::Mat outImage(input_size, input_size, CV_16FC1, output_mems[0]->virt_addr);
        outImage.convertTo(outImage_, CV_8U , 200 / soure_score); // 将像素值缩放到0-255范围
        double minVal, maxVal;
        //cv::minMaxLoc(image, &minVal, &maxVal);
        cv::minMaxLoc(outImage_, &minVal, &maxVal);
        LogDebug << "outImage_ : "<< maxVal;
        // //计算外框1

        // cv::Mat binaryImage;
        // cv::threshold(outImage_, binaryImage, objThresh * (255.0 / 100.0), 255 , cv::THRESH_BINARY);
        // std::vector<std::vector<cv::Point>> contours;
        // cv::findContours(binaryImage, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

        // float scaleX = static_cast<float>(srcSize.width) / input_size;
        // float scaleY = static_cast<float>(srcSize.height) / input_size;
        // od_results->count = 0;
        // for (size_t i = 0; i < contours.size() && od_results->count < OBJ_NUMB_MAX_SIZE; ++i) {
        //     cv::Rect boundingRect = cv::boundingRect(contours[i]);
        //     image_rect_t rect;
        //     rect.left = static_cast<int>(boundingRect.x * scaleX);
        //     rect.top = static_cast<int>(boundingRect.y * scaleY);
        //     rect.right = static_cast<int>((boundingRect.x + boundingRect.width) * scaleX);
        //     rect.bottom = static_cast<int>((boundingRect.y + boundingRect.height) * scaleY);

        //     object_detect_result result;
        //     result.box = rect;
        //     result.prop = pred_score_; // 可以根据需求设置置信度
        //     result.cls_id = 0; // 类别ID, 可以根据需求设置

        //     od_results->results[od_results->count++] = result;
        // }

        // 找到最小值和最大值
        // double min_val, max_val;
        // cv::minMaxLoc(float_image, &min_val, &max_val);

        // // 映射到 [0, 255]
        // cv::Mat normalized_image;
        // float_image.convertTo(normalized_image, CV_8U, 255.0 / (max_val - min_val), -min_val * 255.0 / (max_val - min_val));

        // // 保存图像
        // cv::imwrite(filename, normalized_image);



    }
    else{
        LogWarning << "PatchCore modelDetal['type'].toString() != 'FP16'" ;
        return -1;
    }
    return 0;
}

int PatchCoreAi::pre_training()
{
    LogDebug << "Attempting to start Python Env";
    py::initialize_interpreter();  // 初始化 Python 解释器
    py::module sys = py::module::import("sys");
    //sys.attr("executable") = python_executable;
    py::print(sys.attr("executable"));
    py::print(sys.attr("path"));

    pyTrain = py::module::import("master.scripts.Packaging");
    pyTrainClass = pyTrain.attr("MyModel")();  // 创建 Python 类的实例

    LogDebug << "Attempting to PatchCoreAi::pre_training";
    int ret = pyTrainClass.attr("initial")().cast<int>();
    if(ret!=0)
    {
        ErrorDialogManager::instance().showNonBlockingError("预训练失败", "无法初始化预训练环境");
        LogDebug << "Fail to PatchCoreAi::pre_training :" << ret;
        return -1;
    }
    LogDebug << "successful to PatchCoreAi::pre_training";
    return 0;
}


int PatchCoreAi::load_training_data()
{
    LogDebug << "Attempting to load_training_data: " << modelBeloneDir.path().toStdString();
    int ret = pyTrainClass.attr("load_training_data")(modelBeloneDir.path().toStdString()).cast<int>();
    if(ret!=0)
    {
        ErrorDialogManager::instance().showNonBlockingError("载入图片数据失败", "检查环境");
        LogDebug << "Load training data failed with error code::" << ret;
        return -1;
    }
    LogDebug << "successful to load_training_data";
    return 0;
}

int PatchCoreAi::start_training()
{
    LogDebug << "Attempting to start_training : " << modelDetal["coreSet"].toInt();
    if(modelDetal["coreSet"].toInt() < 10)
    {
        LogError << "start_training failed with coreSetDetal['coreSet'].toInt() < 10 ";
        ErrorDialogManager::instance().showNonBlockingError("开始训练失败", "请先载入基础模型");
        return -1;
    }
    int ret = pyTrainClass.attr("start_training")(modelDetal["coreSet"].toInt()).cast<int>();
    if(ret!=0)
    {
        ErrorDialogManager::instance().showNonBlockingError("开始训练失败", "检查环境");
        LogError << "start_training failed with error code::" << ret;
        return -1;
    }
    LogDebug << "successful to start_training";
    return 0;
}

int PatchCoreAi::continue_iteration()
{
    LogDebug << "Attempting to continue_iteration";
    int ret = pyTrainClass.attr("continue_iteration")().cast<int>();
    if(ret == -1 || ret == -2)
    {
        ErrorDialogManager::instance().showNonBlockingError("开始训练失败", "检查环境");
        LogDebug << "continue_iteration failed with error code::" << ret;
        return -1;
    }else if(ret == -8)
    {
        currentSchedule = 1.1;
        LogDebug << "finish Train" << ret;
        return 0;
    }
    currentSchedule = (float)ret/modelDetal["coreSet"].toInt();
    LogDebug << "successful to continue_iteration" << currentSchedule;
    return 0;
}

void PatchCoreAi::quit_training()
{
    py::finalize_interpreter();  // 终止 Python 解释器
}



void PatchCoreAi::addTrainImage(const cv::Mat img)
{
    if (img.empty()){
        ErrorDialogManager::instance().showNonBlockingError("添加训练集失败", "请先预览或者点击单次");
        LogError << "Attempting to AIprocess::addTrainImage()";
        return;
    }
    // Increment the image index
    imageIndex++;

    // Generate the filename
    QString filename = QString("%1_D.jpg").arg(imageIndex, 3, 10, QChar('0'));

    // Construct the full path
    QString fullPath = modelBeloneDir.filePath(filename);

    // Convert QString to std::string for OpenCV
    std::string path = fullPath.toStdString();

    // Save the image
    if (cv::imwrite(path, img)) {
        LogDebug << "Image saved successfully: " << path;
    } else {
        LogError << "Failed to save image: " << path ;
        ErrorDialogManager::instance().showNonBlockingError("保存图片失败", "保存图片失败");
    }

    // Log the operation
    LogDebug <<  " Added training image: " << filename.toStdString()
             << " (index: " << imageIndex << ")" ;

    QVariantMap imageInfo;
    imageInfo["istrain"] = false;  // 根据需求设置是否已训练
    imageInfo["img"] = fullPath;
    imageInfoList.append(imageInfo);
}

void PatchCoreAi::delTrainImage(int index)
{
    if (index >= 0 && index < imageInfoList.size()) {
        // 获取要删除的元素
        QVariantMap element = imageInfoList[index].toMap();
        QString imagePath = element.value("img").toString();

        // 打印调试信息
        qDebug() << "Removing item at index:" << index << "with image path:" << imagePath;

        // 删除图片文件
        QFile file(imagePath);
        if (file.exists()) {
            if (file.remove()) {
                qDebug() << "Successfully deleted image:" << imagePath;
            } else {
                qDebug() << "Failed to delete image:" << imagePath;
            }
        } else {
            qDebug() << "Image file does not exist:" << imagePath;
        }

        // 从列表中移除元素
        imageInfoList.removeAt(index);
    } else {
        qDebug() << "Index out of bounds:" << index;
    }
}

void PatchCoreAi::clearTrainImage()
{
    if (!modelBeloneDir.exists()) {
        qDebug() << "Directory does not exist:" << modelBeloneDir.path();
        return;
    }

    QStringList files = modelBeloneDir.entryList(QDir::Files);
    foreach (QString file, files) {
        QString filePath = modelBeloneDir.absoluteFilePath(file);
        QFile fileObject(filePath);
        if (fileObject.remove()) {
            qDebug() << "Successfully deleted file:" << filePath;
        } else {
            qDebug() << "Failed to delete file:" << filePath;
        }
    }

    qDebug() << "All files in directory" << modelBeloneDir.path() << "have been deleted.";
    imageInfoList.clear();
}
